from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json"):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        print("state: ", state)
        print("goal: ", goal)
        print("landscape: ", landscape)
    return landscape, state, goal

# save current landscape as json file
def save_current_task(landscape, start_x, start_y, goal_x, goal_y, pth = "./logs/landscape.json"):
    landscape_ = []
    for j in range(landscape[0].shape[0]):
        landscape_.append(int(landscape[0][j]))

    with open(pth, "w") as f:
        json.dump({"data": landscape_, 
                   "state": [start_x, start_y],
                   "goal": [goal_x, goal_y]}, f)


def rnn_run(x, mat_intr, bias1):
    intr_vector = np.dot(x, mat_intr) + bias1
    intr_vector = np.tanh(intr_vector)
    return intr_vector

# 计算雅克比矩阵
@jax.jit
def rnn_run_vector(x, mat_intr, bias1):
    intr_vector = jnp.dot(x, mat_intr) + bias1
    intr_vector = jnp.tanh(intr_vector)
    return intr_vector - x
jacobian_fun = jax.jacrev(rnn_run_vector)
jacobian_vmap = jax.jit(jax.vmap(jacobian_fun, in_axes=(0,None,None)))

@jax.jit
def vector_field_taylor_expansion(x, x0, jacobian, mat_intr, bias1):
    vector = jnp.dot(jacobian, x - x0) + rnn_run_vector(x0, mat_intr, bias1)
    return vector
vector_field_taylor_expansion_vmap = jax.jit(jax.vmap(vector_field_taylor_expansion, in_axes=(0,0,None,None,None)))

# run linear dynamical system using vector_field_taylor_expansion
def run_LDS_one_step(x, x0, dt, jacobian, mat_intr, bias1):

    x_dot = vector_field_taylor_expansion(x, x0, jacobian, mat_intr, bias1)
    x1 = x + x_dot * dt

    return x1.copy()

def compute_linearized_err_1X(sample_point, mat_intr, bias1):

    jacobian_of_x = jacobian_fun(jnp.array(sample_point), mat_intr, bias1)

    # 计算 jacobian_of_x 的所有特征值
    eigenvalues = np.linalg.eigvals(jacobian_of_x)
    # 对 eigenvalues 按照实部从大到小进行排序
    eigenvalues = np.sort(eigenvalues)[::-1]
    print("sorted eigenvalues of jacobian_of_x: ", eigenvalues)

    # 判断 eigenvalues 中是否含有实部大于零的元素
    if any(np.real(eigenvalues) > 0):
        print("存在实部大于零的项")
    else:
        print("不存在实部大于零的项")

    # 将 eigenvalues 所有实部绘制成 barchart
    fig, ax = plt.subplots()
    ax.bar(np.arange(len(eigenvalues)), np.real(eigenvalues))
    plt.show()

    sample_vector = rnn_run_vector(sample_point, mat_intr, bias1)
    sample_vector_len = np.linalg.norm(sample_vector)

    initial_state = sample_point
    current_state = initial_state.copy()
    current_state0 = current_state.copy()
    traj_len = 0

    next_state = rnn_run(initial_state, mat_intr, bias1)

    # while traj_len < sample_vector_len:
    for i in range(2000):
        current_state0 = current_state.copy()
        current_state = run_LDS_one_step(current_state, initial_state, 0.01, jacobian_of_x, mat_intr, bias1)
        delta_state = current_state - initial_state
        delta_state_len = np.linalg.norm(delta_state)
        traj_len += delta_state_len
        current_state0 = current_state.copy()

    diff_next_state_current_state = np.linalg.norm(next_state - current_state0)
    pears_corr = np.corrcoef(next_state, current_state0)[0, 1]

    current_state0_dot = vector_field_taylor_expansion(current_state0, initial_state, jacobian_of_x, mat_intr, bias1)
    next_state_vector = rnn_run_vector(next_state, mat_intr, bias1)

    pears_corr_vector = np.corrcoef(current_state0_dot, next_state_vector)[0, 1]

    print("pears_corr_vector: ", pears_corr_vector)

    return diff_next_state_current_state, pears_corr, current_state0, next_state


def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)
    parser.add_argument("--load_data", type=int, default=rpl_config.load_data)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i
    rpl_config.load_data = args.load_data

    """ create landscape
    """
    random_task = True
    # check if file on rpl_config.task_pth exists
    if os.path.isfile(rpl_config.task_pth):
        random_task = False

    if random_task:
        landscape = generate_maze_pool(num_mazes=1, width=10, height=10)
        landscape = padding_landscapes(landscape, width=12, height=12)
    else:
        landscape, state, goal = load_task(pth = rpl_config.task_pth)
        landscape = [landscape]

    """ create grid env
    """
    start_time = time.time()
    GE = GridEnv(landscapes = landscape, width = 12, height = 12, num_envs_per_landscape = 1, reward_free=True)
    GE.reset()
    print("time taken to create envs: ", time.time() - start_time)

    # 采集所有位置的 general_obs
    # general_obs 是形状为 (12, 12, 10) 的 numpy 数组
    general_obs = np.zeros((12, 12, 10))
    for ii in range(1,11):      # width
        for jj in range(1,11):  # height
            batched_states = jnp.array([[jj,ii] for _ in range(1)])
            concat_obs = get_ideal_obs_vmap_rf(GE.batched_envs, batched_states, GE.batched_goals, GE.last_batched_goal_reached)
            general_obs[jj, ii] = concat_obs[0]

    if not random_task:
        # set states of GE
        GE.batched_states = GE.batched_states.at[0, 0].set(state[0])
        GE.batched_states = GE.batched_states.at[0, 1].set(state[1])
        # set goals of GE
        GE.batched_goals = GE.batched_goals.at[0, 0].set(goal[0])
        GE.batched_goals = GE.batched_goals.at[0, 1].set(goal[1])
        GE.init_batched_states, GE.init_batched_goals = jnp.copy(GE.batched_states), jnp.copy(GE.batched_goals)
        GE.batched_goal_reached = batch_compute_goal_reached(GE.batched_states, GE.batched_goals)
        GE.last_batched_goal_reached = jnp.copy(GE.batched_goal_reached)
        GE.concat_obs = get_ideal_obs_vmap_rf(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)

    concat_obs = GE.concat_obs

    print("concat_obs: ", concat_obs.shape)

    action_list = np.array([[0, 1], [0, -1], [1, 0], [-1, 0], [0, 0]])

    arrow_length = 10
    arrow_list = np.array([[arrow_length, 0], [-arrow_length, 0], [0, arrow_length], [0, -arrow_length], [0, 0]])

    GE.batched_states = jnp.array([[2,4] for i in range(1)])

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    
    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])
    mat1 = np.array(tree_leaves[1])
    print("mat1.shape: ", mat1.shape)
    print("bias1: ", bias1)

    mat_obs = np.array(tree_leaves[1])[rpl_config.nn_size:rpl_config.nn_size+10,:]
    mat_intr = np.array(tree_leaves[1])[0:rpl_config.nn_size,:]
    print("mat_obs.shape: ", mat_obs.shape)

    mat_action = np.array(tree_leaves[3])
    print("mat_action.shape: ", mat_action.shape)

    rnn_state_trajectory_np = np.load("./logs/rnn_state_trajectory_np.npy")

    # 对 rnn_state_np 进行 PCA
    intrinsic_pca = PCA()
    intrinsic_pca.fit(rnn_state_trajectory_np)
    rnn_state_trajectory_np_pca = intrinsic_pca.transform(rnn_state_trajectory_np)

    GE.concat_obs = get_ideal_obs_vmap_rf(GE.batched_envs, GE.batched_states, GE.batched_goals, GE.last_batched_goal_reached)

    def action_distribution(rnn_state, GE, general_obs):
        img = GE.render2()
        # draw actions
        for ii in range(1,11):      # width
            for jj in range(1,11):  # height
                
                IPF = np.dot(general_obs[jj,ii], mat_obs)
                new_state = np.dot(rnn_state, mat_intr) + bias1
                new_state = np.tanh(new_state + IPF)
                action_output = np.dot(new_state, mat_action)
                action = np.argmax(action_output)

                # 检查 action 方向是否有障碍物
                pix = int(GE.batched_envs[0, jj+action_list[action][0], ii+action_list[action][1]])
                pix0 = int(GE.batched_envs[0, jj, ii])
                if pix == 1 and pix0 == 1:
                    start_point = (int(ii)*20+10, int(jj)*20+10)
                    end_point = (start_point[0] + arrow_list[action][0], start_point[1] + arrow_list[action][1])
                    if action == 0:
                        img = cv2.arrowedLine(img, start_point, end_point, (0, 0, 155), 2, tipLength=0.7)
                    elif action == 1:
                        img = cv2.arrowedLine(img, start_point, end_point, (0, 155, 0), 2, tipLength=0.7)
                    elif action == 2:
                        img = cv2.arrowedLine(img, start_point, end_point, (155, 0, 0), 2, tipLength=0.7)
                    elif action == 3:
                        img = cv2.arrowedLine(img, start_point, end_point, (155, 155, 0), 2, tipLength=0.7)
        return img

    while True:

        img = GE.render2()

        # 产生一个 0-rnn_state_trajectory_np.shape[0] 的随机整数
        random_index = random.randint(0, rnn_state_trajectory_np.shape[0]-1)
        action_output = np.dot(rnn_state_trajectory_np[random_index], mat_action)
        action = np.argmax(action_output)
        print("action: ", action)
        print("random_index: ", random_index)

        # draw actions
        action_distribution_img = action_distribution(rnn_state_trajectory_np[random_index], GE, general_obs)

        dim0 = 0
        dim1 = 1
        dim2 = 2

        # 将图像从 BGR 转换为 RGB 格式
        img = cv2.cvtColor(action_distribution_img, cv2.COLOR_BGR2RGB)

        # 创建包含两个子图的图表
        fig = plt.figure(figsize=(15, 5))
        ax1 = fig.add_subplot(121, projection='3d')
        ax2 = fig.add_subplot(122)

        # 在第一个子图中显示您的三维图表
        ax1.scatter(np.array(rnn_state_trajectory_np_pca)[:, dim0], 
                    np.array(rnn_state_trajectory_np_pca)[:, dim1], 
                    np.array(rnn_state_trajectory_np_pca)[:, dim2],
                    s=10)
        
        # 在 random_index 处绘制一个大红点
        ax1.scatter(np.array(rnn_state_trajectory_np_pca)[random_index, dim0], 
                    np.array(rnn_state_trajectory_np_pca)[random_index, dim1], 
                    np.array(rnn_state_trajectory_np_pca)[random_index, dim2],
                    s=300, c="red")

        # 在第二个子图中显示您的 OpenCV 图像
        ax2.imshow(img)

        # # 显示图表
        plt.show()



if __name__ == "__main__":

    ivf()